package cn.getech.data.development.utils;

import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.hive.ql.lib.*;
import org.apache.hadoop.hive.ql.parse.*;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class HiveLineageInfo implements NodeProcessor {

    private static final String sqlRegex = "(?<=[a-zA-Z])\\s(?=[a-zA-Z])";

    /**
     * Stores input tables in sql.
     */
    Set<String> inputTableList = new HashSet();
    /**
     * Stores output tables in sql.
     */
    Set<String> outputTableList = new HashSet();

    Set<String> dropTableList = new HashSet<>();

    /**
     *
     * @return java.util.TreeSet
     */
    public Set<String> getInputTableList() {
        return inputTableList;
    }

    /**
     * @return java.util.TreeSet
     */
    public Set<String> getOutputTableList() {
        return outputTableList;
    }

    public Set<String> getDropTableList() {
        return dropTableList;
    }

    /**
     * Implements the process method for the NodeProcessor interface.
     */
    public Object process(Node nd, Stack stack, NodeProcessorCtx procCtx,
                          Object... nodeOutputs) {
        ASTNode pt = (ASTNode) nd;

        switch (pt.getToken().getType()) {

            case HiveParser.TOK_CREATETABLE:
            case HiveParser.TOK_TAB:
                outputTableList.add(BaseSemanticAnalyzer.getUnescapedName((ASTNode)pt.getChild(0)));
                break;

            case HiveParser.TOK_TABREF:
                ASTNode tabTree = (ASTNode) pt.getChild(0);
                String table_name = (tabTree.getChildCount() == 1) ?
                        BaseSemanticAnalyzer.getUnescapedName((ASTNode)tabTree.getChild(0)) :
                        BaseSemanticAnalyzer.getUnescapedName((ASTNode)tabTree.getChild(0)) + "." + tabTree.getChild(1);
                inputTableList.add(table_name);
                break;
            case HiveParser.TOK_DROPTABLE:
                ASTNode tabTree1 = (ASTNode) pt.getChild(0);
                String table_name1 = (tabTree1.getChildCount() == 1) ?
                        BaseSemanticAnalyzer.getUnescapedName((ASTNode)tabTree1.getChild(0)) :
                        BaseSemanticAnalyzer.getUnescapedName((ASTNode)tabTree1.getChild(0)) + "." + tabTree1.getChild(1);
                dropTableList.add(table_name1);
                break;
        }
        return null;
    }

    /**
     * parses given query and gets the lineage info.
     *
     * @param query
     * @throws ParseException
     */
    public void getLineageInfo(String query) throws ParseException,
            SemanticException {

        /*
         * Get the AST tree
         */
        ParseDriver pd = new ParseDriver();
        ASTNode tree = pd.parse(query);

        while ((tree.getToken() == null) && (tree.getChildCount() > 0)) {
            tree = (ASTNode) tree.getChild(0);
        }

        /*
         * initialize Event Processor and dispatcher.
         */
        inputTableList.clear();
        outputTableList.clear();
        dropTableList.clear();

        // create a walker which walks the tree in a DFS manner while maintaining
        // the operator stack. The dispatcher
        // generates the plan from the operator tree
        Map<Rule, NodeProcessor> rules = new LinkedHashMap<>();

        // The dispatcher fires the processor corresponding to the closest matching
        // rule and passes the context along
        Dispatcher disp = new DefaultRuleDispatcher(this, rules, null);
        GraphWalker ogw = new DefaultGraphWalker(disp);

        // Create a list of topop nodes
        ArrayList topNodes = new ArrayList();
        topNodes.add(tree);
        ogw.startWalking(topNodes, null);
    }

    public String changeHiveSql(String olderSql){
        if(StringUtils.isEmpty(olderSql)){
            return "";
        }
        final String subst = "";
        final Pattern pattern = Pattern.compile(sqlRegex, Pattern.MULTILINE);
        final Matcher matcher = pattern.matcher(olderSql);

        // The substituted value will be contained in the result variable
        return matcher.replaceAll(subst);
    }



    public static void main(String[] args) throws ParseException,
            SemanticException {
        String query = "INSERT OVERWRITE TABLE test.user_stat select name,count(1) as count  from test.user_detail group by name";
        String sql1 = "select * from ods.jjj";
        String sql = "drop table ods.jjj";
        HiveLineageInfo lep = new HiveLineageInfo();
        lep.getLineageInfo(query);
        System.out.println("Input tables = " + lep.getInputTableList());
        System.out.println("Output tables = " + lep.getOutputTableList());
        System.out.println("drop tables = " + lep.getDropTableList());


        final String string = "ab cd 1234567 1234567 ef gh 1234567 1234567 ij kl - - - -";
        final String subst = "";

        final Pattern pattern = Pattern.compile(sqlRegex, Pattern.MULTILINE);
        final Matcher matcher = pattern.matcher(string);

        // The substituted value will be contained in the result variable
        final String result = matcher.replaceAll(subst);

        System.out.println("Substitution result: " + result);
    }
}